"""
Description:
Author: 月间
Date: 2024-11-24 17:19:16
LastEditTime: 2024-11-24 21:22:38
LastEditors:
"""

import json

from fastapi import Depends, Query, APIRouter
from fastapi_utils.cbv import cbv
from funboost import ApsJobAdder, PublisherParams, BoostersManager, RedisMixin
from funboost.constant import RedisKeys
from funboost.core.active_cousumer_info_getter import (
    QueueConusmerParamsGetter,
    ActiveCousumerProcessInfoGetter,
)
from funboost.function_result_web.functions import Statistic
from funboost.timing_job.apscheduler_use_redis_store import (
    FunboostBackgroundSchedulerProcessJobsWithinRedisLock,
)
from motor.motor_asyncio import AsyncIOMotorDatabase, AsyncIOMotorClient
from pyecharts import options as opts
from pyecharts.charts import Bar
from pyecharts.faker import Faker
from sqlalchemy import and_

from authx.dependencies.permission import PermissionChecker
from src import schemas, models
from src.core.base_controller import BaseController
from src.core.dependencies import IdList
from src.core.response.response_schema import (
    ResponseGenericSchema,
    ResponseSchema,
    PageResponse,
)
from src.db.db_database import alchemy
from src.db.db_mongodb import get_mongo_database
from src.log import logger
from src.schemas.funcboost_record_schema import FuncBoostRecordSchema
from src.schemas.task_schema import HeartbeatInfo, QueueOut, TaskParams
from src.services.task_service import TaskService

router = APIRouter()


@cbv(router)
class TaskController(BaseController):
    task_service: TaskService = Depends(alchemy.provide_service(TaskService))
    """
    定时任务接口
    """

    @router.get(
        "/list",
        summary="查询定时任务列表",
        response_model=PageResponse[schemas.TaskOut],
        dependencies=[Depends(PermissionChecker("system:task:list"))],
    )
    async def get_task_list(self, params: TaskParams = Depends()):
        """
        查询定时任务列表
        :return:
        """
        where = await params.get_where_and_limit_offset(models.SysTask)
        res, count = await self.task_service.list_and_count(*where)
        return self.paginated_response(
            data=res, total=count, current=params.page, size=params.limit
        )

    @router.post(
        "/add",
        summary="添加定时任务",
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:task:add"))],
    )
    async def add_task(self, data: schemas.TaskCreate):
        """
        添加定时任务
        :return:
        """
        # async with self.db.begin():
        task = await self.task_service.create(data, auto_commit=True)
        logger.warning(f"添加任务成功：{task.to_dict()}")
        await self.task_service.save_task(task)
        return self.success("添加成功")

    @router.put(
        "/update",
        summary="更新定时任务",
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:task:update"))],
    )
    async def update_task(self, data: schemas.TaskUpdate):
        """
        更新定时任务
        :return:
        """
        task = await self.task_service.update(data, auto_commit=True, auto_refresh=True)
        await self.task_service.save_task(task)
        return self.success("更新成功")

    @router.put(
        "/change/status",
        summary="暂停/恢复定时任务",
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:task:update"))],
    )
    async def change_status(self, data: schemas.TaskActive):
        """
        暂停/恢复定时任务
        :return:
        """
        task = await self.task_service.update(data, auto_commit=True, auto_refresh=True)
        # 判断任务是否存在
        aps: FunboostBackgroundSchedulerProcessJobsWithinRedisLock = (
            ApsJobAdder.get_funboost_redis_apscheduler(task.name)
        )
        job = aps.get_job(str(data.id))
        if not job:
            await self.task_service.save_task(task)
        else:
            if data.is_active:
                # 恢复
                aps.resume_job(str(data.id))
            else:
                # 暂停
                aps.pause_job(str(data.id))
        return self.success("更新成功")

    @router.delete(
        "/delete",
        summary="删除定时任务",
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:task:delete"))],
    )
    async def delete_task(self, item: IdList = Depends()):
        """
        删除定时任务
        :return:
        """
        # 查询所有删除任务
        jobs = await self.task_service.list(and_(models.SysTask.id.in_(item.ids)))
        await self.task_service.delete_many(item.ids)
        for job in jobs:
            aps: FunboostBackgroundSchedulerProcessJobsWithinRedisLock = (
                ApsJobAdder.get_funboost_redis_apscheduler(job.name)
            )
            aps.resume_job(str(job.id))
        return self.success("删除成功")

    @router.get(
        "/job/all",
        summary="获取所有任务列表",
    )
    async def get_all_tasks_list(
        self,
        col_name_search: str | None = Query(None, description="模糊搜索任务名"),
        mongo: AsyncIOMotorClient = Depends(get_mongo_database("task_status")),
    ):
        """
        获取所有任务列表
        :return:
        """
        if not col_name_search:
            collection_name_list = await mongo.list_collection_names()
        else:
            collection_name_list = [
                collection_name
                for collection_name in await mongo.list_collection_names()
                if col_name_search in collection_name
            ]
        res = [
            {
                "label": collection_name,
                "value": collection_name,
            }
            for collection_name in collection_name_list
        ]
        return self.success(data=res)

    @router.get(
        "/record/list",
        summary="获取所有任务日志",
        response_model=PageResponse[FuncBoostRecordSchema],
    )
    async def get_all_tasks_log(
        self,
        queue_name: str = Query(None, description="任务名"),
        current: int = Query(1, description="页码"),
        size: int = Query(10, description="每页数量"),
        success: bool = Query(None, description="是否成功"),
        engine: AsyncIOMotorDatabase = Depends(get_mongo_database("task_status")),
    ):
        """
        获取所有任务日志
        :return:
        """
        # 判断task_name是否存在
        if not queue_name:
            return self.success(
                data={"total": 0, "records": [], "current": current, "size": size}
            )
        filter_dict = {}
        if success is not None:
            filter_dict = {"success": success}
        # 时间倒序
        res = (
            await engine.get_collection(queue_name)
            .find(filter=filter_dict)
            .sort("insert_time_str", -1)
            .skip((current - 1) * size)
            .limit(size)
            .to_list(length=size)
        )
        total = await engine.get_collection(queue_name).count_documents(filter_dict)
        return self.paginated_response(
            data=res,
            total=total,
            current=current,
            size=size,
        )

    @router.get(
        "/queue/statistics",
        summary="统计：最近60秒最近60分钟，最近24小时，最近10天执行情况",
    )
    async def get_all_tasks_statistics(
        self,
        task_name: str = Query(..., description="任务名"),
        time_type: str = Query(..., description="时间类型"),
    ):
        """
        60分钟，最近24小时，最近7天执行情况
        :return:
        """
        time_dict = {
            "recent_10_days": "最近10天消费情况",
            "recent_24_hours": "最近24小时消费情况",
            "recent_60_minutes": "最近60分钟消费情况",
            "recent_60_seconds": "最近60秒消费情况",
        }
        if time_type not in time_dict:
            return self.error(message="时间类型错误")
        if not task_name:
            return self.success()
        stat = Statistic(task_name)
        stat.build_result()
        data = stat.result.get(time_type)

        c = (
            Bar()
            .add_xaxis(
                data.get("time_arr", Faker.choose()),
            )
            .add_yaxis(task_name, data.get("count_arr", Faker.values()))
            .set_global_opts(
                xaxis_opts=opts.AxisOpts(axislabel_opts=opts.LabelOpts(rotate=-15)),
                title_opts=opts.TitleOpts(title=time_dict.get(task_name)),
                datazoom_opts=opts.DataZoomOpts(),
                toolbox_opts=opts.ToolboxOpts(),
                legend_opts=opts.LegendOpts(is_show=False),
            )
        )
        return self.success(data=json.loads(c.dump_options()))

    # 获取所有任务队列列表
    @router.get(
        "/queue/list",
        summary="获取所有任务队列列表",
        tags=["任务管理"],
        response_model=PageResponse[QueueOut],
    )
    async def get_all_tasks_queue(
        self,
    ):
        data = QueueConusmerParamsGetter().get_queue_params_and_active_consumers()
        result = []
        for k, v in data.items():
            result.append(
                {
                    "queue_name": k,
                    "broker_kind": v.get("queue_params", {"broker_kind": ""}).get(
                        "broker_kind"
                    ),
                    # 函数累计平均耗时
                    "total_avg_time": v.get(
                        "all_consumers_avarage_function_spend_time_from_start", 0
                    ),
                    # 最近10秒的平均耗时
                    "last_ten_seconds_time": v.get(
                        "all_consumers_last_x_s_avarage_function_spend_time", 0
                    ),
                    # 最近10秒的运行完成次数
                    "last_ten_seconds_num": v.get(
                        "all_consumers_last_x_s_execute_count", 0
                    ),
                    # 最近10秒的运行失败次数
                    "last_ten_seconds_fail_num": v.get(
                        "all_consumers_last_x_s_execute_count_fail", 0
                    ),
                    # 消费者数量
                    "consumer_num": len(v.get("active_consumers", [])),
                    # 是否暂停
                    "pause_flag": v.get("pause_flag"),
                }
            )
        return self.paginated_response(data=result, total=len(result))

    # 查询队列中的消息数量
    @router.get(
        "/queue/msg/num",
        summary="查询队列中的消息数量",
        tags=["任务管理"],
        response_model=ResponseGenericSchema[int],
    )
    async def get_queue_msg_num(
        self,
        broker_kind: str = Query(..., description="broker类型"),
        queue_name: str = Query(..., description="队列名"),
    ):
        publisher = BoostersManager.get_cross_project_publisher(
            PublisherParams(
                queue_name=queue_name,
                broker_kind=broker_kind,
                publish_msg_log_use_full_msg=True,
            )
        )
        return self.success(data=publisher.get_message_count())

    @router.post(
        "/queue/clear",
        summary="清空队列中的消息",
        tags=["任务管理"],
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:queue:clear"))],
    )
    async def clear_queue(
        self,
        broker_kind: str = Query(..., description="broker类型"),
        queue_name: str = Query(..., description="队列名"),
    ):
        publisher = BoostersManager.get_cross_project_publisher(
            PublisherParams(
                queue_name=queue_name,
                broker_kind=broker_kind,
                publish_msg_log_use_full_msg=True,
            )
        )
        publisher.clear()
        return self.success()

    @router.post(
        "/queue/pause",
        summary="暂停队列",
        tags=["任务管理"],
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:queue:pause"))],
    )
    async def pause_queue(
        self,
        queue_name: str = Query(..., description="队列名"),
    ):
        RedisMixin().redis_db_frame.hset(
            RedisKeys.REDIS_KEY_PAUSE_FLAG, queue_name, "1"
        )
        return self.success()

    @router.post(
        "/queue/resume",
        summary="恢复队列",
        tags=["任务管理"],
        response_model=ResponseSchema,
        dependencies=[Depends(PermissionChecker("system:queue:resume"))],
    )
    async def resume_queue(
        self,
        queue_name: str = Query(..., description="队列名"),
    ):
        RedisMixin().redis_db_frame.hset(
            RedisKeys.REDIS_KEY_PAUSE_FLAG, queue_name, "0"
        )
        return self.success()

    @router.get(
        "/consumers",
        summary="获取运行中的消费者",
        tags=["任务管理"],
        response_model=ResponseGenericSchema[list[dict]],
    )
    async def get_consumers(
        self,
        partition_by: str = Query(..., description="分区方式，可选值为'ip'或'queue'"),
    ):
        if partition_by == "ip":
            info_map = ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_partition_by_ip()
        elif partition_by == "queue":
            info_map = ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_partition_by_queue_name()
        else:
            return self.error(message="无效的分区方式，可选值为'ip'或'queue'")

        ret_list = []
        total_count = 0
        for k, v in info_map.items():
            ret_list.append({"collection_name": k, "count": len(v)})
            total_count += len(v)
        ret_list = sorted(ret_list, key=lambda x: x["collection_name"])
        ret_list.insert(0, {"collection_name": "所有", "count": total_count})
        return self.success(data=ret_list)

    @router.get(
        "/consumers/options",
        summary="获取运行中的消费者",
        tags=["任务管理"],
        response_model=ResponseGenericSchema[list[dict]],
    )
    async def get_consumers_options(
        self,
    ):
        info_map = ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_partition_by_queue_name()
        ret_list = []
        for k, v in info_map.items():
            ret_list.append({"label": k, "value": k})
        return self.success(data=ret_list)

    @router.get(
        "/heartbeat/queue",
        summary="根据队列名获取心跳信息",
        tags=["任务管理"],
        response_model=PageResponse[HeartbeatInfo],
    )
    async def heartbeat_info_by_queue_name(
        self,
        queue_name: str = Query(None, description="队列名，可选值为'所有'或具体队列名"),
    ):
        if queue_name in ("所有", None):
            info_map = ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_partition_by_queue_name()
            ret_list = []
            for queue_name, dic in info_map.items():
                ret_list.extend(dic)
            return self.paginated_response(data=ret_list, total=len(ret_list))
        else:
            data = (
                ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_by_queue_name(
                    queue_name
                )
            )
            return self.paginated_response(data=data, total=len(data))

    @router.get(
        "/heartbeat/ip",
        summary="根据IP获取心跳信息",
        tags=["任务管理"],
        response_model=PageResponse[HeartbeatInfo],
    )
    async def heartbeat_info_by_ip(
        self,
        ip: str = Query(None, description="IP地址，可选值为'所有'或具体IP"),
    ):
        if ip in ("所有", None):
            info_map = ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_partition_by_ip()
            ret_list = []
            for ip, dic in info_map.items():
                ret_list.extend(dic)
            return self.success(data=ret_list)
        else:
            return self.success(
                data=ActiveCousumerProcessInfoGetter().get_all_hearbeat_info_by_ip(ip)
            )
